import numpy as np

import matplotlib.pyplot as plt

from util.util_ris_pattern import point_2_phase, phase_2_bit, phaseBit_2_deg, phase_2_pattern_xyz_fft

from multi_beam_trace.util_create_trace import create_scan_traces_phi, \
    create_scan_traces_theta, create_scan_traces_theta_phi
from multi_beam_trace.util_phase_pattern import traces_2_phaseBits
from multi_beam_trace.util_analysis import compare_arrays


#================================================== 流程相关 ==========================================================

#----------------------------------------------------- 统计不一样的位数 ------------------------------------------------------------
# 保存轨迹变化
def save_plots(list_data1, list_data2, list_data3, list_str, path_img_pre):
    # 检查列表长度是否一致
    if not (len(list_data1) == len(list_data2) == len(list_str)):
        raise ValueError("All input lists must have the same length.")
    # 逐图保存
    i = 1
    for i, (data1, data2, data3, str_label) in enumerate(zip(list_data1, list_data2, list_data3, list_str)):
        # 创建一个新的图像
        fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(9, 3))
        # 在左上方绘制 list_data1[i]
        im1 = ax1.imshow(data1, aspect='auto')
        ax1.set_title(f"{str_label[0]}")
        ax1.axis('equal')
        ax1.axis('off')
        # plt.colorbar(im1, ax=ax1)     # 注释掉能关闭图例
        # 在中上方绘制 list_data2[i]
        im2 = ax2.imshow(data2, aspect='auto')
        ax2.set_title(f"{str_label[1]}")
        ax2.axis('equal')
        ax2.axis('off')
        # plt.colorbar(im2, ax=ax2)
        # 在右上方绘制 list_data3[i]
        im3 = ax3.imshow(data3, aspect='auto')
        ax3.set_title(f"{str_label[2]}")
        ax3.axis('equal')
        ax3.axis('off')
        # plt.colorbar(im3, ax=ax3)
        # 在下方放置文本
        fig.text(0.5, 0.01, str_label[3], ha='center', fontsize=12)
        # 调整子图间距
        plt.subplots_adjust(wspace=0.1, bottom=0.05)
        # 保存图像
        # plt.show()
        path_img = path_img_pre + str(i) + ".jpg"
        plt.savefig(path_img)
        # 关闭当前图像，释放内存
        plt.close(fig)
        i += 1

# 统计phases相邻phase[i]不相同的位数
def count_diff(phaseBits):
    counts = []
    for i in range(0, len(phaseBits)-1):
        phaseBit1 = phaseBits[i]
        phaseBit2 = phaseBits[i+1]
        # 比较 phase1 和 phase2 的每一位，值不相同则 count 加 1
        counts.append(np.sum(phaseBit1 != phaseBit2))
    return counts

# 每幅图保存: 相邻位置 - diff - count(diff)
def save_phases_diff(traces, phaseBits, counts, path_dir):
    list_data1 = []
    list_data2 = []
    list_data3 = []
    list_str = []
    for i in range(len(phaseBits) - 1):
        phase_this = phaseBits[i]
        phase_next = phaseBits[i + 1]
        diff_arr = compare_arrays(phase_this, phase_next)
        img_strs = [str(traces[i]), str(traces[i + 1]), "diff", "count diff:" + str(counts[i]) + ", total: 4096"]
        list_data1.append(phase_this)
        list_data2.append(phase_next)
        list_data3.append(diff_arr)
        list_str.append(img_strs)
    save_plots(list_data1, list_data2, list_data3, list_str, path_dir)

# 每幅图保存: phase - pattern - diff - count(diff)
def save_phase_pattern_diff(bit_num, traces, phaseBits, counts, path_dir):
    list_data1 = []
    list_data2 = []
    list_data3 = []
    list_str = []
    for i in range(len(phaseBits)):
        phaseBit = phaseBits[i]
        # fft计算方向图
        phaseDeg = phaseBit_2_deg(phaseBit, bit_num)
        phaseRad = np.deg2rad(phaseDeg)
        # pattern = phase_2_pattern(phaseBit_mix)                # 公式法直接计算, 准确但速度太慢
        pattern, x, y = phase_2_pattern_xyz_fft(phaseRad)  # FFT法计算, 快速
        # parse
        if i > 0:
            phaseBit_next = phaseBits[i-1]
            diff_arr = compare_arrays(phaseBit, phaseBit_next)
            img_strs = [str(traces[i]), "pattern", "diff", "count diff:" + str(counts[i - 1]) + ", total: 4096"]
        else:
            rows, cols = phaseBit.shape
            diff_arr = np.zeros((rows, cols))
            img_strs = [str(traces[i]), "pattern", "diff", ""]
        list_data1.append(phaseBit)
        list_data2.append(pattern)
        list_data3.append(diff_arr)
        list_str.append(img_strs)
    save_plots(list_data1, list_data2, list_data3, list_str, path_dir)


# =================================================== 主函数 ======================================================
# 沿theta方向统计变化
def trace_counts_theta(theta_start, theta_end, phi, path_pre):
    bit_num = 1
    traces = create_scan_traces_theta(theta_start, theta_end, phi)
    phaseBits = traces_2_phaseBits(traces, bit_num)
    counts = count_diff(phaseBits)
    print("counts:", counts)
    # 每幅图保存: phase(t=i), phase(t=i+1), diff[phase(t=i+1) - phase(t=i)], count(diff)
    save_phases_diff(traces, phaseBits, counts, path_pre + "compare_phases_diff/")
    # 每幅图保存: phase(t=i), pattern(t=i+1), diff[phase(t=i) - phase(t=i-1)], count(diff)
    save_phase_pattern_diff(bit_num, traces, phaseBits, counts, path_pre + "compare_phase_pattern_diff/")


# 沿phi方向统计变化
def trace_counts_phi(theta, phi_start, phi_end, path_pre):
    bit_num = 1
    traces = create_scan_traces_phi(theta, phi_start, phi_end)
    phaseBits = traces_2_phaseBits(traces, bit_num)
    counts = count_diff(phaseBits)
    print("counts:", counts)
    # 每幅图保存: phase(t=i), phase(t=i+1), diff[phase(t=i+1) - phase(t=i)], count(diff)
    save_phases_diff(traces, phaseBits, counts, path_pre + "compare_phases_diff/")
    # 每幅图保存: phase(t=i), pattern(t=i+1), diff[phase(t=i) - phase(t=i-1)], count(diff)
    save_phase_pattern_diff(bit_num, traces, phaseBits, counts, path_pre + "compare_phase_pattern_diff/")


# 沿(theta,phi)方向统计变化
def trace_counts_theta_phi(theta_start, phi_start, theta_phi_length, path_pre):
    bit_num = 1
    traces = create_scan_traces_theta_phi(theta_start, phi_start, theta_phi_length)
    phaseBits = traces_2_phaseBits(traces, bit_num)
    counts = count_diff(phaseBits)
    print("counts:", counts)
    # 每幅图保存: phase(t=i), phase(t=i+1), diff[phase(t=i+1) - phase(t=i)], count(diff)
    save_phases_diff(traces, phaseBits, counts, path_pre + "compare_phases_diff/")
    # 每幅图保存: phase(t=i), pattern(t=i+1), diff[phase(t=i) - phase(t=i-1)], count(diff)
    save_phase_pattern_diff(bit_num, traces, phaseBits, counts, path_pre + "compare_phase_pattern_diff/")


def main():
    # trace_counts_theta(0, 90, 0, "../files/multi-beam-trace/beam1/theta(0,1,90)-phi(0,0,0)/")
    # trace_counts_theta(0, 90, 15, "../files/multi-beam-trace/beam1/theta(0,1,90)-phi(15,0,15)/")
    # trace_counts_theta(0, 90, 30, "../files/multi-beam-trace/beam1/theta(0,1,90)-phi(30,0,30)/")
    # trace_counts_theta(0, 90, 45, "../files/multi-beam-trace/beam1/theta(0,1,90)-phi(45,0,45)/")
    # trace_counts_theta(0, 90, 60, "../files/multi-beam-trace/beam1/theta(0,1,90)-phi(60,0,60)/")
    # trace_counts_theta(0, 90, 90, "../files/multi-beam-trace/beam1/theta(0,1,90)-phi(90,0,90)/")
    # trace_counts_theta(0, 90, 180, "../files/multi-beam-trace/beam1/theta(0,1,90)-phi(180,0,180)/")
    # trace_counts_theta(0, 90, 270, "../files/multi-beam-trace/beam1/theta(0,1,90)-phi(270,0,270)/")
    #
    # trace_counts_phi(10, 0, 360, "../files/multi-beam-trace/beam1/theta(10,0,10)-phi(0,1,360)/")
    # trace_counts_phi(15, 0, 360, "../files/multi-beam-trace/beam1/theta(15,0,15)-phi(0,1,360)/")
    # trace_counts_phi(30, 0, 360, "../files/multi-beam-trace/beam1/theta(30,0,30)-phi(0,1,360)/")
    # trace_counts_phi(45, 0, 360, "../files/multi-beam-trace/beam1/theta(45,0,45)-phi(0,1,360)/")
    # trace_counts_phi(60, 0, 360, "../files/multi-beam-trace/beam1/theta(60,0,60)-phi(0,1,360)/")
    #
    # trace_counts_theta_phi(10, 10, 60, "../files/multi-beam-trace/beam1/theta(10,1,70)-phi(10,1,70)/")
    trace_counts_theta_phi(10, 90, 60, "../files/multi-beam-trace/beam1/theta(10,1,70)-phi(90,1,150)/")
    trace_counts_theta_phi(10, 180, 60, "../files/multi-beam-trace/beam1/theta(10,1,70)-phi(180,1,240)/")




if __name__=="__main__":
    main()